from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker, declarative_base
from app.core.config import settings
import logging
import clickhouse_driver
import asyncio
from aiocache import Cache, cached
from aiocache.serializers import JsonSerializer
from contextlib import asynccontextmanager

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 创建 Base 类
Base = declarative_base()

# PostgreSQL 数据库 URL
POSTGRES_DATABASE_URL = f"postgresql+asyncpg://{settings.POSTGRES_USER}:{settings.POSTGRES_PASSWORD}@{settings.POSTGRES_HOST}:{settings.POSTGRES_PORT}/{settings.POSTGRES_DB}"
logger.info(f"PostgreSQL 连接: {settings.POSTGRES_HOST}:{settings.POSTGRES_PORT}/{settings.POSTGRES_DB}")

# CitusDB 数据库 URL (与 PostgreSQL 使用相同的驱动，但连接到不同的实例)
CITUSDB_DATABASE_URL = f"postgresql+asyncpg://{settings.CITUSDB_USER}:{settings.CITUSDB_PASSWORD}@{settings.CITUSDB_HOST}:{settings.CITUSDB_PORT}/{settings.CITUSDB_DB}"
logger.info(f"CitusDB 连接: {settings.CITUSDB_HOST}:{settings.CITUSDB_PORT}/{settings.CITUSDB_DB}")

# 创建 PostgreSQL 异步引擎
postgres_engine = create_async_engine(
    POSTGRES_DATABASE_URL,
    echo=settings.SQL_ECHO,
    future=True,
    pool_size=settings.DB_POOL_SIZE,
    max_overflow=settings.DB_MAX_OVERFLOW,
    pool_timeout=settings.DB_POOL_TIMEOUT,
    pool_recycle=settings.DB_POOL_RECYCLE
)

# 创建 CitusDB 异步引擎
citusdb_engine = create_async_engine(
    CITUSDB_DATABASE_URL,
    echo=settings.SQL_ECHO,
    future=True,
    pool_size=settings.DB_POOL_SIZE,
    max_overflow=settings.DB_MAX_OVERFLOW,
    pool_timeout=settings.DB_POOL_TIMEOUT,
    pool_recycle=settings.DB_POOL_RECYCLE
)

# 创建 PostgreSQL 异步会话工厂
PostgresSessionLocal = sessionmaker(
    postgres_engine,
    class_=AsyncSession,
    expire_on_commit=False,
    autoflush=False
)

# 创建 CitusDB 异步会话工厂
CitusDBSessionLocal = sessionmaker(
    citusdb_engine,
    class_=AsyncSession,
    expire_on_commit=False,
    autoflush=False
)

# ClickHouse 客户端连接池
clickhouse_clients = {}

# 缓存配置
cache_config = {
    'ttl': settings.CACHE_TTL,
    'serializer': JsonSerializer()
}

@asynccontextmanager
async def get_clickhouse_client():
    """获取 ClickHouse 客户端"""
    global clickhouse_clients
    current_task = asyncio.current_task()
    if current_task not in clickhouse_clients:
        clickhouse_clients[current_task] = clickhouse_driver.Client(
            host=settings.CLICKHOUSE_HOST,
            port=settings.CLICKHOUSE_PORT,
            user=settings.CLICKHOUSE_USER,
            password=settings.CLICKHOUSE_PASSWORD,
            database=settings.CLICKHOUSE_DB,
            settings={'use_numpy': True}
        )
    try:
        yield clickhouse_clients[current_task]
    finally:
        pass

@asynccontextmanager
async def get_postgres_db():
    """获取 PostgreSQL 数据库会话的上下文管理器"""
    async with PostgresSessionLocal() as session:
        try:
            yield session
        finally:
            await session.close()

@asynccontextmanager
async def get_citusdb():
    """获取 CitusDB 数据库会话的上下文管理器"""
    async with CitusDBSessionLocal() as session:
        try:
            yield session
        finally:
            await session.close()

async def get_db():
    """获取 PostgreSQL 数据库会话的依赖函数"""
    async with PostgresSessionLocal() as session:
        try:
            yield session
        finally:
            await session.close()

async def get_citus_db():
    """获取 CitusDB 数据库会话的依赖函数"""
    async with CitusDBSessionLocal() as session:
        try:
            yield session
        finally:
            await session.close() 